#ifndef ARMOR_DETECTOR_NUMBER_CLASSIFIER_H
#define ARMOR_DETECTOR_NUMBER_CLASSIFIER_H

#include <opencv2/opencv.hpp>
#include <vector>
#include <string>
#include <mutex>

#include "armor_detector_types.h"

class NumberClassifier {
public:
    NumberClassifier(const std::string &model_path,
                     const std::string &label_path,
                     const double threshold,
                     const std::vector<std::string> &ignore_classes = {});

    // Extract the roi image of number from the src
    cv::Mat extractNumber(const cv::Mat &src, const Armor &armor) const;

    // Classify the number of the armor
    void classify(const cv::Mat &src, Armor &armor);

    // Erase the ignore classes
    void eraseIgnoreClasses(std::vector<Armor> &armors);

    double threshold;

private:
    std::mutex mutex_;
    cv::dnn::Net net_;
    std::vector<std::string> class_names_;
    std::vector<std::string> ignore_classes_;
};

#endif // ARMOR_DETECTOR_NUMBER_CLASSIFIER_H
